#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
树莓派部署脚本
用于在树莓派上部署和运行FunASR语音识别模型
"""

import os
import argparse
import time
import threading
import queue
import numpy as np
import sounddevice as sd
import onnxruntime as ort
import librosa
import matplotlib.pyplot as plt

# 配置参数
parser = argparse.ArgumentParser(description='树莓派语音识别部署脚本')
parser.add_argument('--model_path', type=str, default='./models/paraformer_int8.onnx', help='量化模型路径')
parser.add_argument('--sample_rate', type=int, default=16000, help='采样率')
parser.add_argument('--duration', type=int, default=2, help='每次录音时长(秒)')
parser.add_argument('--output_dir', type=str, default='./output', help='输出目录')
args = parser.parse_args()

# 创建输出目录
os.makedirs(args.output_dir, exist_ok=True)

class ASRService:
    def __init__(self, model_path, sample_rate=16000):
        """
        初始化语音识别服务
        
        参数:
            model_path: ONNX模型路径
            sample_rate: 采样率
        """
        # 配置ONNX Runtime会话选项
        options = ort.SessionOptions()
        options.intra_op_num_threads = 3  # 使用3个线程进行并行计算
        options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        options.enable_mem_pattern = True
        options.enable_mem_reuse = True
        
        # 加载ONNX模型
        self.session = ort.InferenceSession(model_path, options)
        self.sample_rate = sample_rate
        self.audio_queue = queue.Queue()
        self.result_queue = queue.Queue()
        self.is_running = False
        
        # 字符映射表（简化版）
        self.id_to_char = {}
        for i in range(8404):
            self.id_to_char[i] = chr(i + 0x4E00) if i < 3500 else chr((i % 3500) + 0x4E00)
        self.id_to_char[0] = ''  # blank
        
    def start(self):
        """启动语音识别服务"""
        self.is_running = True
        self.recognition_thread = threading.Thread(target=self._recognition_worker)
        self.recognition_thread.daemon = True
        self.recognition_thread.start()
        print("语音识别服务已启动")
        
    def stop(self):
        """停止语音识别服务"""
        self.is_running = False
        if hasattr(self, 'recognition_thread'):
            self.recognition_thread.join(timeout=1.0)
        print("语音识别服务已停止")
            
    def _recognition_worker(self):
        """语音识别工作线程"""
        while self.is_running:
            try:
                audio_data = self.audio_queue.get(timeout=0.5)
                
                # 提取特征
                features = self._extract_features(audio_data)
                
                # 执行推理
                start_time = time.time()
                input_name = self.session.get_inputs()[0].name
                result = self.session.run(None, {input_name: features})[0]
                end_time = time.time()
                
                # 解码结果
                text = self._decode_result(result)
                
                # 将结果放入队列
                inference_time = end_time - start_time
                self.result_queue.put((text, inference_time))
                self.audio_queue.task_done()
                
            except queue.Empty:
                continue
            except Exception as e:
                print(f"识别过程中出错: {str(e)}")
                
    def _extract_features(self, audio_data):
        """从音频数据提取特征"""
        try:
            # 提取梅尔频谱特征
            mel_spec = librosa.feature.melspectrogram(
                y=audio_data, 
                sr=self.sample_rate,
                n_fft=400,
                hop_length=160,
                win_length=400,
                n_mels=80,
                fmin=0,
                fmax=8000
            )
            
            # 转换为对数刻度
            log_mel_spec = librosa.power_to_db(mel_spec)
            
            # 归一化
            mean = np.mean(log_mel_spec)
            std = np.std(log_mel_spec)
            normalized_features = (log_mel_spec - mean) / (std + 1e-5)
            
            # 添加批次维度
            features = np.expand_dims(normalized_features, axis=0).astype(np.float32)
            return features
            
        except Exception as e:
            print(f"特征提取出错: {str(e)}")
            # 返回空特征
            return np.zeros((1, 80, 500), dtype=np.float32)
        
    def _decode_result(self, result):
        """解码模型输出为文本"""
        try:
            # 获取最可能的字符ID序列
            ids = np.argmax(result, axis=1)
            
            # 移除重复的ID和空白标记
            prev_id = -1
            text_ids = []
            for id in ids:
                if id != prev_id and id != 0:  # 不等于前一个ID且不是空白
                    text_ids.append(id)
                prev_id = id
            
            # 转换为文本
            text = ''.join([self.id_to_char.get(id, '') for id in text_ids])
            return text
            
        except Exception as e:
            print(f"解码出错: {str(e)}")
            return ""
        
    def process_audio(self, audio_data):
        """处理音频数据"""
        self.audio_queue.put(audio_data)
        
    def get_result(self, timeout=None):
        """获取识别结果"""
        try:
            return self.result_queue.get(timeout=timeout)
        except queue.Empty:
            return None, 0

def record_audio(duration, sample_rate):
    """
    录制音频
    
    参数:
        duration: 录制时长(秒)
        sample_rate: 采样率
        
    返回:
        音频数据
    """
    print(f"录制{duration}秒音频...")
    audio_data = sd.rec(
        int(duration * sample_rate),
        samplerate=sample_rate,
        channels=1,
        dtype='float32'
    )
    sd.wait()  # 等待录制完成
    return audio_data.flatten()

def monitor_system_resources():
    """
    监控系统资源使用情况
    """
    import psutil
    
    # 获取当前进程
    process = psutil.Process(os.getpid())
    
    # 记录资源使用情况
    timestamps = []
    cpu_usages = []
    memory_usages = []
    
    start_time = time.time()
    
    try:
        while True:
            # 记录时间戳
            current_time = time.time() - start_time
            timestamps.append(current_time)
            
            # 记录CPU使用率
            cpu_percent = psutil.cpu_percent(interval=0.1)
            cpu_usages.append(cpu_percent)
            
            # 记录内存使用情况
            memory_info = process.memory_info()
            memory_mb = memory_info.rss / (1024 * 1024)  # 转换为MB
            memory_usages.append(memory_mb)
            
            # 打印当前资源使用情况
            print(f"时间: {current_time:.1f}s, CPU: {cpu_percent}%, 内存: {memory_mb:.1f}MB")
            
            # 每秒更新一次
            time.sleep(1)
            
            # 如果运行超过60秒，则停止监控
            if current_time > 60:
                break
                
    except KeyboardInterrupt:
        pass
    
    # 绘制资源使用情况图表
    plt.figure(figsize=(12, 6))
    
    plt.subplot(1, 2, 1)
    plt.plot(timestamps, cpu_usages)
    plt.xlabel('时间 (秒)')
    plt.ylabel('CPU使用率 (%)')
    plt.title('CPU使用率随时间变化')
    plt.grid(True)
    
    plt.subplot(1, 2, 2)
    plt.plot(timestamps, memory_usages)
    plt.xlabel('时间 (秒)')
    plt.ylabel('内存使用 (MB)')
    plt.title('内存使用随时间变化')
    plt.grid(True)
    
    plt.tight_layout()
    plt.savefig(os.path.join(args.output_dir, 'resource_usage.png'))
    plt.close()
    
    # 计算平均资源使用情况
    avg_cpu = sum(cpu_usages) / len(cpu_usages)
    avg_memory = sum(memory_usages) / len(memory_usages)
    max_memory = max(memory_usages)
    
    print(f"平均CPU使用率: {avg_cpu:.1f}%")
    print(f"平均内存使用: {avg_memory:.1f}MB")
    print(f"最大内存使用: {max_memory:.1f}MB")
    
    return {
        'avg_cpu': avg_cpu,
        'avg_memory': avg_memory,
        'max_memory': max_memory
    }

def compare_models_on_raspberry_pi():
    """
    比较不同模型在树莓派上的性能
    """
    # 模拟数据
    models = ['WeNet (INT8)', 'Whisper-small', 'FunASR (INT8)']
    memory_usage = [320, 850, 180]  # MB
    cpu_usage = [85, 98, 65]  # %
    inference_times = [3.5, 68, 2.2]  # 秒
    
    # 创建内存使用对比图
    plt.figure(figsize=(10, 6))
    bars = plt.bar(models, memory_usage, color=['skyblue', 'salmon', 'lightgreen'])
    
    # 添加数值标签
    for bar in bars:
        height = bar.get_height()
        plt.text(bar.get_x() + bar.get_width()/2., height + 10,
                f'{height} MB', ha='center', va='bottom')
    
    plt.title('树莓派上不同模型的内存占用对比')
    plt.ylabel('内存占用 (MB)')
    plt.grid(axis='y', linestyle='--', alpha=0.7)
    plt.tight_layout()
    plt.savefig(os.path.join(args.output_dir, 'memory_usage_comparison.png'))
    plt.close()
    
    # 创建CPU使用率对比图
    plt.figure(figsize=(10, 6))
    bars = plt.bar(models, cpu_usage, color=['skyblue', 'salmon', 'lightgreen'])
    
    # 添加数值标签
    for bar in bars:
        height = bar.get_height()
        plt.text(bar.get_x() + bar.get_width()/2., height + 1,
                f'{height}%', ha='center', va='bottom')
    
    plt.title('树莓派上不同模型的CPU使用率对比')
    plt.ylabel('CPU使用率 (%)')
    plt.ylim(0, 105)
    plt.grid(axis='y', linestyle='--', alpha=0.7)
    plt.tight_layout()
    plt.savefig(os.path.join(args.output_dir, 'cpu_usage_comparison.png'))
    plt.close()
    
    # 创建推理时间对比图
    plt.figure(figsize=(10, 6))
    bars = plt.bar(models, inference_times, color=['skyblue', 'salmon', 'lightgreen'])
    
    # 添加数值标签
    for bar in bars:
        height = bar.get_height()
        plt.text(bar.get_x() + bar.get_width()/2., height + 1,
                f'{height}秒', ha='center', va='bottom')
    
    plt.title('树莓派上不同模型处理2秒语音的推理时间对比')
    plt.ylabel('推理时间 (秒)')
    plt.grid(axis='y', linestyle='--', alpha=0.7)
    plt.tight_layout()
    plt.savefig(os.path.join(args.output_dir, 'inference_time_comparison.png'))
    plt.close()
    
    # 创建树莓派硬件规格图表
    specs = {
        'CPU': '4核 ARM Cortex-A72 @ 1.5GHz',
        '内存': '4GB LPDDR4',
        '存储': '32GB microSD卡',
        '操作系统': 'Raspberry Pi OS (64位)',
        '电源': '5V/3A USB-C',
        '接口': 'USB 3.0, HDMI, 以太网, Wi-Fi, 蓝牙'
    }
    
    fig, ax = plt.subplots(figsize=(10, 6))
    ax.axis('tight')
    ax.axis('off')
    
    table_data = [[k, v] for k, v in specs.items()]
    
    table = ax.table(cellText=table_data,
                    colLabels=['硬件组件', '规格'],
                    cellLoc='center',
                    loc='center')
    
    table.auto_set_font_size(False)
    table.set_fontsize(12)
    table.scale(1, 1.5)
    
    plt.title('树莓派4B硬件规格')
    plt.tight_layout()
    plt.savefig(os.path.join(args.output_dir, 'hardware_specs.png'))
    plt.close()

def main():
    # 比较不同模型在树莓派上的性能
    compare_models_on_raspberry_pi()
    
    # 初始化语音识别服务
    asr_service = ASRService(args.model_path, args.sample_rate)
    asr_service.start()
    
    # 启动资源监控线程
    monitor_thread = threading.Thread(target=monitor_system_resources)
    monitor_thread.daemon = True
    monitor_thread.start()
    
    try:
        # 测试循环
        for i in range(5):
            print(f"\n测试 {i+1}/5")
            
            # 录制音频
            audio_data = record_audio(args.duration, args.sample_rate)
            
            # 处理音频
            asr_service.process_audio(audio_data)
            
            # 获取识别结果
            text, inference_time = asr_service.get_result(timeout=10.0)
            
            if text:
                print(f"识别结果: {text}")
                print(f"推理时间: {inference_time:.2f}秒")
            else:
                print("识别超时或失败")
            
            time.sleep(1)  # 短暂暂停
            
    except KeyboardInterrupt:
        print("\n用户中断，正在退出...")
    finally:
        # 停止服务
        asr_service.stop()
        
        print("测试完成")

if __name__ == "__main__":
    main()
